#!/usr/bin/env python

from Bio.PDB import *
import csv
import re
import os

#import dbtbl
import xml.etree.ElementTree as ET
import urllib

#uniprot XML parser
def GetUniprotFile(proteinID):
    if type(proteinID) == type(set()) or type(proteinID) == type(list()):
        for iter in proteinID:  
            uniprot_file = urllib.urlretrieve('http://www.uniprot.org/uniprot/'+iter+'.xml', 'uniprot/'+iter+'.xml')
    elif type(proteinID) == type(str()):
        uniprot_file = urllib.urlretrieve('http://www.uniprot.org/uniprot/'+proteinID+'.xml', 'uniprot/'+proteinID+'.xml')
    else:
        raise ("Unkonw proteinID type.")
    
def GetName(protein, ns = '{http://uniprot.org/uniprot}'):
    '''Get the name of the protein'''
    try:
        name = (protein.find(ns+'protein')
            .find(ns+'recommendedName')
                .findtext(ns+'fullName'))
    except:
        name = (protein.find(ns+'protein')
                .find(ns+'submittedName')
                .findtext(ns+'fullName'))
        
    return name
    
def GetAccessions(protein, ns = '{http://uniprot.org/uniprot}'):
    '''Get all accession codes for this protein'''
    accession_list = []
    accessions = protein.findall(ns+'accession')
    for accession in accessions:
        accession_list.append(accession.text)
    return accession_list 

def GetGeneName(protein, ns = '{http://uniprot.org/uniprot}'):
    '''Get the gene name'''
    try:
        gene_name = protein.find(ns+'gene').findtext(ns+'name')
    except:
        # Have to invent short names for proteins that don't have one
        try:
            fullname = (protein.find(ns+'protein')
                        .find(ns+'recommendedName')
                        .findtext(ns+'fullName'))
        except:
            fullname = (protein.find(ns+'protein')
                        .find(ns+'submittedName')
                        .findtext(ns+'fullName'))
        gene_name = fullname[0:13]
    gene_synonym = []
    try:
        names = protein.find(ns+'gene').findall(ns+'name')
    except:
        names = []
    if len(names) > 1:
        
        for name in names:
            if name.attrib['type'] == 'synonym':
                gene_synonym.append(name.text)
    else:
        names = ''
    return gene_name, gene_synonym 

def GetSequence(protein, ns = '{http://uniprot.org/uniprot}'):
    '''Retrieve the protein sequence'''
    sequence = protein.findtext(ns+'sequence')
    seq_clean = re.sub("\s+", "", sequence)
    #seq_wrapped = textwrap.fill(seq_clean, 60)
    return seq_clean

def GetOrganism(protein, ns = '{http://uniprot.org/uniprot}'): 
    '''Retrieve the protein organism'''
    organism_name = 'None'
    organism_list = protein.find(ns+'organism').findall(ns+'name')
    for organism in organism_list:  
        if organism.get('type') == 'scientific':
            organism_name = organism.text
    taxon_list = protein.find(ns+'organism').findall(ns+'dbReference')
    for taxon in taxon_list:
        taxonid = int(taxon.get('id'))
    return organism_name, taxonid

def GetLocation(protein, ns = '{http://uniprot.org/uniprot}'):
    '''return length with start and stop'''
    
    l_begin = (protein.find(ns+'feature').
               find(ns+'location').
               find(ns+'begin').
               get('position'))
    l_end = (protein.find(ns+'feature')
             .find(ns+'location')
             .find(ns+'end')
             .get('position'))
    return int(l_begin), int(l_end)
    

def GetDomain(protein, ns = '{http://uniprot.org/uniprot}'):
    '''Return domains with start and stop positions'''
    domain_dict = {}
    domain_list = []
    domains = protein.findall(ns+'feature')
    # This loop seems cumbersome
    for feature in domains:
        if feature.get('type') == 'domain':
            domain_list.append(feature)
            for domain in domain_list:
                domain_start = (feature.find(ns+'location')
                                .find(ns+'begin')
                                .get('position'))
                domain_stop = (feature.find(ns+'location')
                               .find(ns+'end')
                               .get('position'))
                domain_span = (int(domain_start), int(domain_stop))
                domain_dict[feature.get('description')]= domain_span
    return domain_dict

def GetLength(protein, ns = '{http://uniprot.org/uniprot}'):
    '''return the sequence length of protein'''
    length = protein.find(ns+'sequence').get('length')
    return int(length)

    
def ID2tree(proteinID, ns = '{http://uniprot.org/uniprot}'):
    #input a uniprot accession
    #output a protein
    if os.path.exists('uniprot/'+proteinID+'.xml'):
        pass
    else:
        print "Downloading uniprot file" + proteinID + "...", 
        GetUniprotFile(proteinID)
        print "DONE"
    
    if os.path.getsize('uniprot/'+proteinID+'.xml') != 0:
        #=0 means this entry has already deleted by uniprot
        #tree = ET.parse('/home/sdai/workspace/PDZ/uniprot/'+proteinID+'.xml')
        tree = ET.parse('uniprot/'+proteinID+'.xml')
        for item in tree.getiterator(ns+'uniprot'):
            #global protein
            protein = item.findall(ns+'entry')[0]
        return protein
    else:
        return False

def pdz2_ext(proteinID):
    #if protein got only one pdz domain return false
    #more than one, return domian num, from, to

    protein = ID2tree(proteinID)
    if protein is not False:
        pdz = GetDomain(protein)
        pdz2 = {}
        for k in pdz:
            if k.startswith('PDZ'):
                pdz2[k] = pdz[k]
        
        if len(pdz2) > 0:
            return pdz2
        else:
            print proteinID, "has one or less PDZ domain."
            return False
    else:
        print "Fail to find Uniprot XML file for ", proteinID
        return False
    

def proteininfo(proteinID):
    #add the information for the new tables: protein, taxon and pdz2
    #gene_name
    #taxon_ID
    #seq
    #organism
    #pdz2 
    protein = ID2tree(proteinID)
    if protein is not False:
        proteinname = GetName(protein)
        genename = GetGeneName(protein)
        sequence = GetSequence(protein)
        organism, taxonid = GetOrganism(protein)
        return genename, sequence, organism, taxonid, proteinname
        
    else:    
        return False

def uniprot_ext(proteinID, BindingSiteS, BindingSiteE, ns = '{http://uniprot.org/uniprot}'):
    #
    #return proteinID, BindingSiteS, BindingSiteE, sequence
    #----------------------------------
    
    nameA = '-'
    lenA = '-'
    accessionsA = '-'
    genenameA = '-'
    organismA = '-'
    domain_dictA = '-'
    sequenceA = '-'
    #------------------------------------
    #if os.path.getsize(
    #    '/home/sdai/workspace/PDZ/uniprot/'+proteinID+'.xml') != 0:

    protein = ID2tree(proteinID)
    if protein is not False:
        nameA = GetName(protein)
        lenA = GetLength(protein)
        accessionsA = GetAccessions(protein)
        #genenameA, genesynonymA = GetGeneName(protein)
        #organismA = GetOrganism(protein)
        #domain_dictA = GetDomain(protein)
        sequenceA = GetSequence(protein)
        '''
        if abs(BindingSiteS-BindingSiteE) > 100 or (BindingSiteS == 0 and
                                                    (BindingSiteE > 100
                                                     or BindingSiteE ==0)):
                
                print proteinID, BindingSiteS, BindingSiteE
        '''

        if BindingSiteS == 0 and BindingSiteE == 0:
            return (proteinID, BindingSiteS, BindingSiteE, sequenceA)
            
        elif BindingSiteE == 0:
            return (proteinID, BindingSiteS, BindingSiteE,
                    sequenceA[BindingSiteS:])
        elif BindingSiteS ==0:
            return (proteinID, BindingSiteS, BindingSiteE,
                    sequenceA[:BindingSiteE])
        elif BindingSiteE-BindingSiteS <= 5:
            return (proteinID, BindingSiteS, BindingSiteE,
                    sequenceA[BindingSiteE-5:BindingSiteE])
        else:
            return (proteinID, BindingSiteS, BindingSiteE,
                    sequenceA[BindingSiteS:BindingSiteE])
    else:
        return False, False, False, False

def load_peptide(filename):
    # load peptide1.out generated by peptide1.py
    pepdb = dict()
    input_table = csv.reader(open(filename), delimiter=' ')
    for row in input_table:
        pepdb[row[1]] = (row[0], row[2])
    return pepdb

if __name__ == '__main__':
    
   
    ns = '{http://uniprot.org/uniprot}'
    tree = ET.parse('uniprot/'+'Q8VBX6'+'.xml')
    for item in tree.getiterator(ns+'uniprot'):
        #global protein
        protein = item.findall(ns+'entry')[0]
        
    print pdz2_ext('O35889')